import jieba
from rank_bm25 import BM25Okapi

class KnowledgeRetriever:
    def retrieve(self, query,documents,top_k=10):
        # print(documents)
        # 实现知识检索逻辑
        # for doc in documents:
        #     print(list(jieba.cut(doc)))
        tokenized_corpus = [list(jieba.cut(doc)) for doc in documents]
        # 初始化BM25对象
        bm25 = BM25Okapi(tokenized_corpus)
        tokenized_query = list(jieba.cut(query))
        # print(tokenized_query)
        top = bm25.get_top_n(tokenized_query, documents, n=top_k)
        return top


if __name__ == '__main__':
    pass
